
package cn.myeasyai.milvus;

import cn.myeasyai.config.MilvusConfig;
import cn.myeasyai.dto.ImageDTO;
import cn.myeasyai.feature.FeatureExtractor;
import com.alibaba.fastjson.JSON;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.DataType;
import io.milvus.grpc.SearchResults;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import io.milvus.param.R;
import io.milvus.param.RpcStatus;
import io.milvus.param.collection.CreateCollectionParam;
import io.milvus.param.collection.FieldType;
import io.milvus.param.collection.FlushParam;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.InsertParam.Field;
import io.milvus.param.dml.SearchParam;
import io.milvus.param.index.CreateIndexParam;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

import jakarta.annotation.Resource;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.springframework.stereotype.Service;


/**
 * @author fushoujiang
 */
@Service
public class VectorService {
    @Resource
    private MilvusServiceClient milvusClient;

    @Resource
    private MilvusConfig milvusConfig;


    private void createXunjianCollection() {
        FieldType idField = FieldType.newBuilder().withName("id").withDataType(DataType.Int64).withPrimaryKey(true).build();
        FieldType vectorField = FieldType.newBuilder().withName("embedding").withDataType(DataType.FloatVector).withDimension(1000).build();
        FieldType nameField = FieldType.newBuilder().withName("urlName").withDataType(DataType.VarChar).withMaxLength(1000).build();
        CreateCollectionParam createCollectionParam = CreateCollectionParam.newBuilder().withCollectionName("xunjianImage").withDescription("巡检系统标准图片").withShardsNum(2).addFieldType(idField).addFieldType(vectorField).addFieldType(nameField).build();
        this.milvusClient.createCollection(createCollectionParam);
    }

    public void search(List<Float> feature) {
        List<List<Float>> vectors = Collections.singletonList(feature);
        SearchParam searchParam = SearchParam.newBuilder().withCollectionName(milvusConfig.getCollectionName())
                .withMetricType(MetricType.L2).withTopK(5).withVectors(vectors)
                .withVectorFieldName(milvusConfig.getVectorFieldName()).build();
        R<SearchResults> searchResults = this.milvusClient.search(searchParam);
        System.out.println("Searching vector: " + JSON.toJSONString(feature));
        System.out.println("Result: " + JSON.toJSONString(searchResults.toString()));
    }

    /**
     * 初始化向量数据
     *
     * @param imageDTO 用户信息
     * @param features 向量数据
     */
    public void initImageVector(ImageDTO imageDTO, List<Float> features) {
        List<String> url = Collections.singletonList(imageDTO.getUrl());
        List<List<Float>> vectors = Collections.singletonList(features);
        List<Field> fields = new ArrayList();
        fields.add(new Field(milvusConfig.getVectorFieldName(), vectors));
        fields.add(new Field("image_url", url));
        InsertParam insertParam = InsertParam.newBuilder().withCollectionName(milvusConfig.getCollectionName()).withFields(fields).build();
        this.milvusClient.insert(insertParam);
    }


    private void flush(String name) {
        this.milvusClient.flush(FlushParam.newBuilder().withCollectionNames(Collections.singletonList(name)).withSyncFlush(true).withSyncFlushWaitingInterval(50L).withSyncFlushWaitingTimeout(30L).build());
    }

    private void buildIndex(String name) {
        System.out.println("Building AutoIndex...");
        IndexType INDEX_TYPE = IndexType.AUTOINDEX;
        long startIndexTime = System.currentTimeMillis();
        R<RpcStatus> indexR = this.milvusClient.createIndex(CreateIndexParam.newBuilder().withCollectionName(name).withFieldName("embedding").withIndexType(INDEX_TYPE).withMetricType(MetricType.L2).withSyncMode(Boolean.TRUE).withSyncWaitingInterval(500L).withSyncWaitingTimeout(30L).build());
        long endIndexTime = System.currentTimeMillis();
        System.out.println("Succeed in " + (double) (endIndexTime - startIndexTime) / 1000.0D + " seconds!");
    }
}
